#根据预测出来的txt文件裁剪图片
import os
import cv2
from tqdm import tqdm

image_input = '/media/liuyi/024b9406-a087-42e6-8444-51e6a224c28e/Downloads/new_xinnengyuan/'
txt_input = '/media/liuyi/024b9406-a087-42e6-8444-51e6a224c28e/Downloads/vehicle_txt/'
path_output = "/media/liuyi/024b9406-a087-42e6-8444-51e6a224c28e/Downloads/"    # 裁剪出来的小图保存的根目录

img_total = []
txt_total = []

file_image = os.listdir(image_input)
for filename in file_image:
     first,last = os.path.splitext(filename)
     img_total.append(first)

file_txt = os.listdir(txt_input)
for filename in file_txt:
     first,last = os.path.splitext(filename)
     txt_total.append(first)

for img_ in tqdm(img_total):
    if img_ in txt_total:
        filename_img = img_+".jpg"
        path1 = os.path.join(image_input,filename_img)
        img = cv2.imread(path1)
        filename_txt = img_     #预测出来的txt文件没有后缀名，有则加 {+".txt"}
        h = img.shape[0]
        w = img.shape[1]
        n = 1
        with open(os.path.join(txt_input,filename_txt),"r+",encoding="utf-8",errors="ignore") as f:
            for line in f:
                aa = line.split(" ")
                if not int(aa[0]) == 0: continue
                x_center = w * float(aa[1])       # aa[1]左上点的x坐标
                y_center = h * float(aa[2])       # aa[2]左上点的y坐标
                width = int(w*float(aa[3]))       # aa[3]图片width
                height = int(h*float(aa[4]))      # aa[4]图片height
                lefttopx = int(x_center-width/2.0)
                lefttopy = int(y_center-height/2.0)
                roi = img[lefttopy+1:lefttopy+height+3,lefttopx+1:lefttopx+width+1]   # [左上y:右下y,左上x:右下x] (y1:y2,x1:x2)需要调参，否则裁剪出来的小图可能不太好
                if roi.size == 0: continue                   #判断需要裁剪的类别:0--vehicle
                filename_last = img_+"_"+str(n)+".jpg"      # 裁剪出来的小图文件名
                path2 = os.path.join(path_output,"cut_txt")           # 需要在path_output路径下创建一个cut_txt文件夹
                # print('path2:', path2)                    # 裁剪小图的保存位置
                cv2.imwrite(os.path.join(path2,filename_last),roi)
                n = n+1
    else:
        continue
